import pandas as pd
import os
from glob import glob
import shutil

v_info = pd.read_csv('./origin_dataset/dataset.csv').set_index('VideoID')

v_map = pd.read_csv('./origin_dataset/datasetmap.txt', sep='\s+').set_index('VideoID')

video = v_map.join(v_info)

ret = []
for i, row in video.iterrows():
    folder_path = './origin_dataset/dataset_picture/' + row.Video[0:3] + '/' + row.Video[3:] + '/'
    image_paths = glob(folder_path + '*.jpg')
    count_image_in_folder = len(image_paths)
    for image_path in image_paths:
        image_basename = os.path.basename(image_path)
        new_image_path = './data/raw/%s_%s' % (row.Video, image_basename)

        # 处理缺失情况
        delta = 1
        image_id = int(os.path.basename(image_path).split('.')[0])
        while True:
            new_image_id = image_id + delta
            if new_image_id <= count_image_in_folder:
                ref_image_basename = str(new_image_id).zfill(2) + '.jpg'
                image_path = os.path.join(
                    os.path.dirname(image_path),
                    ref_image_basename
                )
                if os.path.getsize(image_path) != 0:
                    break
            new_image_id = image_id - delta
            if new_image_id >= 1:
                ref_image_basename = str(new_image_id).zfill(2) + '.jpg'
                image_path = os.path.join(
                    os.path.dirname(image_path),
                    ref_image_basename
                )
                if os.path.getsize(image_path) != 0:
                    break
            delta += 1

        shutil.copy(image_path, new_image_path)
        ret.append({
            'VideoID': i,
            'ImageID': int(image_basename.split('.')[0]),
            'RawPath': new_image_path
        })

image_video = pd.DataFrame(ret).set_index('VideoID')

all_info = video.join(image_video, how='left')

all_info.to_csv('data/raw.csv')
